package com.btl.service.ai.support.rag.store.vector;

import com.btl.service.ai.support.rag.store.vector.collections.VectorCollectionsEnum;
import dev.langchain4j.model.embedding.EmbeddingModel;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Configuration;

import java.util.Collection;
import java.util.concurrent.ExecutionException;

@Slf4j
@Configuration
public class VectorStoreInit {

    @Resource
    private ApplicationContext applicationContext;

    @Resource
    private QdrantClient qdrantClient;

    @PostConstruct
    public void init() {
        Collection<EmbeddingModel> values = applicationContext
                .getBeansOfType(EmbeddingModel.class)
                .values();
        values.forEach(value -> {
            try {
                Collections.Distance distance = Collections.Distance.Cosine;
                for (VectorCollectionsEnum vectorCollectionsEnum : VectorCollectionsEnum.values()) {
                    String collectionName = vectorCollectionsEnum.getCollection() + "_" + value.dimension();
                    if (!collectionExists(qdrantClient, collectionName)) {
                        qdrantClient.createCollectionAsync(
                                        collectionName,
                                        Collections.VectorParams.newBuilder()
                                                .setDistance(distance)
                                                .setSize(value.dimension())
                                                .build())
                                .get();
                        log.info("Creating Qdrant collection: {}", collectionName);
                    } else {
                        log.info("Collection already exists: {}", collectionName);
                    }
                }
            } catch (InterruptedException | ExecutionException e) {
                throw new RuntimeException(e);
            }
        });
    }

    public boolean collectionExists(QdrantClient qdrantClient, String collectionName) throws ExecutionException, InterruptedException {
        return qdrantClient.collectionExistsAsync(collectionName)
                .get();
    }

}
